import numpy as np
import onnxruntime
import onnx
import argparse
import os

def SaveNumpyArrayToFile(out_data,fileName):
	file = open(fileName, "wb")
	out_data1 = np.ascontiguousarray(out_data,dtype=np.float32)
	print(fileName)
	for x in np.nditer(out_data1):
		file.write(x)
	file.close()

def run_onnx_model(opt):
	inputs={}
	for input in opt.inputs:
		print(input,type(input))
		input_list = input.split("=")
		inputs[input_list[0]] =input_list[1]
	print(inputs)

	input_shapes={}
	for input_shape in opt.input_shapes:
		shape_kv=input_shape.split("=")
		shape_list = []
		print(shape_kv[1])
		shape_str_list = shape_kv[1].split(",")
		for shape_item in shape_str_list:
			print(shape_item)
			shape_list.append(int(shape_item))
		input_shapes[shape_kv[0]] =tuple(shape_list)
	print(input_shapes)

	onnx_model = onnx.load(opt.model_file)

	if opt.IsDumpMidLayer:
		for node in onnx_model.graph.node:
			for out_node in node.output:
				onnx_model.graph.output.extend([onnx.ValueInfoProto(name=out_node)])

	sess = onnxruntime.InferenceSession(onnx_model.SerializeToString(),providers=['CPUExecutionProvider'])
	input_names = sess.get_inputs()
	if len(input_names) != len(inputs):
		print("the input num error. model input num = " , len(input_names), ",len(inputs)=",len(inputs))

	inputs_arg = {}
	for i in range(len(input_names)):
		input_name = input_names[i].name.strip()
		file = open(inputs[input_name], "rb")
		img = file.read()
		img = np.frombuffer(img, dtype=np.float32)
		img = img.reshape(input_shapes[input_name] )
		inputs_arg[input_name] = img
		#print(inputs_arg[input_name].shape)

	output = sess.run(None, inputs_arg)
	if opt.IsSaveOutput:
		all_layer_out = sess.get_outputs()
		for i, out in enumerate(all_layer_out):
			print(out.name, out.type,out.shape)
			tmp_file_name = out.name.replace("/","_") 
			save_file_name = os.path.join(opt.OutputPath,tmp_file_name)
			SaveNumpyArrayToFile(output[i], save_file_name)
	#print(output)
	return output
if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('--model_file', type=str, default='', help='model.pt path(s)')  
	parser.add_argument('--inputs', nargs='+', help='modul multi input name and value, --inputs input1="./input1.bin" input2="./input2.bin" ')  
	parser.add_argument('--input_shapes', nargs='+',  help='modul multi input name and shape, --inputs input1="1,3,640,640" input2="1,3,244,244" ')  
	parser.add_argument('--IsDumpMidLayer',  type=bool ,  default=False, help='Is dump all mid layer?')  
	parser.add_argument('--IsSaveOutput',  type=bool ,  default=False, help='Is dump all mid layer?')  
	parser.add_argument('--OutputPath',  type=str ,  default='./', help='Is dump all mid layer?')
	opt = parser.parse_args()
	print(opt)
	run_onnx_model(opt)

